import random
from faker import Faker
import mysql.connector
from mysql.connector import Error
import time
import sys
import argparse


def generate_room_list():
    """生成房间号列表（共40个房间）"""
    rooms = []
    rooms.extend([str(x) for x in range(101, 111)])  # 10个（101-110）
    rooms.extend([str(x) for x in range(201, 212)])  # 11个（201-211）
    rooms.extend([str(x) for x in range(301, 311)])  # 10个（301-310）
    rooms.extend([str(x) for x in range(501, 510)])  # 9个（501-509）
    return rooms


def generate_dataset(num_people):
    """生成带房间人数限制的数据集"""
    fake = Faker('zh_CN')
    rooms = generate_room_list()
    room_count = len(rooms)

    # 验证人数范围
    min_people = room_count * 10
    max_people = room_count * 30
    if not (min_people <= num_people <= max_people):
        raise ValueError(f"人数必须在{min_people}到{max_people}之间")

    # 初始化每个房间基础人数
    room_distribution = {room: 10 for room in rooms}
    remaining = num_people - min_people

    # 随机分配剩余人数
    while remaining > 0:
        room = random.choice(rooms)
        if room_distribution[room] < 30:
            room_distribution[room] += 1
            remaining -= 1

    # 生成房间分配序列
    room_sequence = []
    for room, count in room_distribution.items():
        room_sequence.extend([room] * count)
    random.shuffle(room_sequence)

    # 生成完整数据集
    dataset = []
    for idx, room in enumerate(room_sequence, 1):
        dataset.append({
            "id": idx,
            "room_number": room,
            "name": fake.name(),
            "join_date": fake.date_between(start_date='-13y', end_date='today').strftime('%Y-%m-%d')
        })

    return dataset


def save_to_mysql(data, host, user, password, database='person_db'):
    """存储数据到MySQL"""
    try:
        print(f"⏳ 正在连接到MySQL数据库 {host}...")
        connection = mysql.connector.connect(
            host=host,
            user=user,
            password=password
        )

        cursor = connection.cursor()

        # 创建数据库（如果不存在）
        cursor.execute("CREATE DATABASE IF NOT EXISTS {} DEFAULT CHARACTER SET 'utf8mb4'".format(database))
        cursor.execute(f"USE {database}")

        # 创建表（如果不存在）
        create_table_query = """
        CREATE TABLE IF NOT EXISTS persons (
            id INT PRIMARY KEY,
            room_number VARCHAR(10),
            name VARCHAR(100),
            join_date DATE
        )
        """
        cursor.execute(create_table_query)

        # 清空现有数据
        cursor.execute("TRUNCATE TABLE persons")

        # 批量插入数据
        print("⏳ 正在插入数据到MySQL...")
        batch_size = 500
        start_time = time.time()

        for i in range(0, len(data), batch_size):
            batch = data[i:i + batch_size]
            insert_query = """
            INSERT INTO persons (id, room_number, name, join_date)
            VALUES (%s, %s, %s, %s)
            """
            records = [(p['id'], p['room_number'], p['name'], p['join_date']) for p in batch]
            cursor.executemany(insert_query, records)
            connection.commit()
            progress = min(i + batch_size, len(data))
            print(f"✅ 已插入 {progress}/{len(data)} 条数据", end='\r')

        cursor.close()
        connection.close()
        # 修复：使用字符串常量而非未定义变量
        print(f"\n✅ 成功存储 {len(data)} 条数据到MySQL表 persons")
        print(f"⏱️ 插入耗时: {time.time() - start_time:.2f}秒")
        return True

    except Error as e:
        print(f"❌ MySQL数据库错误: {str(e)}")
        return False
    except Exception as e:
        print(f"❌ 存储数据到MySQL时出错: {str(e)}")
        return False


def main():
    try:
        # 默认数据库配置
        default_config = {
            "host": "127.0.0.1",
            "user": "root",
            "password": "831114399x",
            "database": "person_db"
        }

        # 尝试从命令行参数获取配置
        parser = argparse.ArgumentParser(description='生成人员数据并存储到MySQL')
        parser.add_argument('--host', help='MySQL主机地址', default=default_config["host"])
        parser.add_argument('--user', help='MySQL用户名', default=default_config["user"])
        parser.add_argument('--password', help='MySQL密码', default=default_config["password"])
        parser.add_argument('--db', help='MySQL数据库名称', default=default_config["database"])
        parser.add_argument('--num', type=int, help='生成的人员数量')

        args = parser.parse_args()

        rooms = generate_room_list()
        min_people = len(rooms) * 10
        max_people = len(rooms) * 30

        # 获取用户输入
        if args.num and min_people <= args.num <= max_people:
            num_people = args.num
        else:
            while True:
                try:
                    num_people = int(input(f"请输入要生成的人数（{min_people}-{max_people}）: "))
                    if min_people <= num_people <= max_people:
                        break
                    print(f"输入值必须介于{min_people}和{max_people}之间")
                except ValueError:
                    print("请输入有效的数字")

        # 生成数据
        start_time = time.time()
        print("⏳ 正在生成数据...")
        dataset = generate_dataset(num_people)
        print(f"✅ 数据生成完成，耗时 {time.time() - start_time:.2f}秒")

        # 存储到MySQL
        start_time = time.time()
        db_config = {
            "host": args.host,
            "user": args.user,
            "password": args.password,
            "database": args.db
        }

        print(f"📂 数据库配置: {db_config['host']}/{db_config['database']} (用户: {db_config['user']})")
        if save_to_mysql(dataset, **db_config):
            print(f"✅ 总耗时: {time.time() - start_time:.2f}秒")

    except KeyboardInterrupt:
        print("\n操作已取消")
        sys.exit(0)
    except Exception as e:
        print(f"❌ 发生错误: {str(e)}")


if __name__ == "__main__":
    main()
